# @file    : llama_index_test
# @time    : 2025/3/22
# @author  : yongpeng.yao
# @desc    :

from dotenv import load_dotenv

load_dotenv()
from llama_index.core import SimpleDirectoryReader

documents = SimpleDirectoryReader('docs').load_data()

# 为文档建立索引
from llama_index.core import VectorStoreIndex

from deepseek_llm import DeepSeekLLM

dashscope_llm = DeepSeekLLM()

index = VectorStoreIndex.from_documents(documents, llm=dashscope_llm)

# 创建查询引擎
agent = index.as_query_engine()

# 两个查询问题
response = agent.query('花语秘境的员工有几种角色？')
print('花语秘境的员工有几种角色？', response)

response = agent.query('花语秘境的Agent叫什么名字？')
print('花语秘境的Agent叫什么名字？', response)

# 把索引保存到本地
index.storage_context.persist()
